package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html/charset"

	"github.com/go-resty/resty/v2"
)

// EastmoneyResponse 东方财富响应结构
type EastmoneyResponse struct {
	Rc     int64                 `json:"rc"`
	Rt     int64                 `json:"rt"`
	Svr    int64                 `json:"svr"`
	Lt     int64                 `json:"lt"`
	Full   int64                 `json:"full"`
	Dlmkts string                `json:"dlmkts"`
	Data   EastmoneyResponseData `json:"data"`
}

// EastmoneyResponseData 东方财富响应数据结构
type EastmoneyResponseData struct {
	Code      string   `json:"code"`
	Market    int64    `json:"market"`
	Name      string   `json:"name"`
	Decimal   int64    `json:"decimal"`
	Dktotal   int64    `json:"dktotal"`
	PreKPrice float64  `json:"preKPrice"`
	Klines    []string `json:"klines"`
}

// IndexData 指数数据结构
type IndexData struct {
	Date       string  `json:"date"`
	ClosePoint float64 `json:"closePoint"`
}

// IndexCode 指数代码结构
type IndexCode struct {
	Code  string `json:"code"`
	Secid string `json:"secid"`
	Name  string `json:"name"`
}

// 全局变量用于存储0条数据的股票代码
var zeroDataCodes []string
var zeroDataMutex sync.Mutex

func main() {
	// 初始化客户端
	client := resty.New().SetTimeout(30 * time.Second)

	// 设置重试机制
	client.SetRetryCount(3)
	client.SetRetryWaitTime(2 * time.Second)
	client.SetRetryMaxWaitTime(5 * time.Second)
	client.AddRetryCondition(func(r *resty.Response, err error) bool {
		return err != nil || r.StatusCode() >= 500
	})

	// 获取指数代码列表
	indexCodeList, err := ListIndexCodes()
	if err != nil {
		log.Fatalf("获取指数代码列表失败: %v", err)
	}

	// 限制并发数
	semaphore := make(chan struct{}, 5) // 限制5个并发请求
	var wg sync.WaitGroup

	// 并发抓取数据
	for _, indexCode := range indexCodeList {
		wg.Add(1)
		semaphore <- struct{}{} // 获取信号量

		go func(code IndexCode) {
			defer wg.Done()
			defer func() { <-semaphore }() // 释放信号量

			if err := FetchData(code, client); err != nil {
				fmt.Printf("抓取数据失败: %v\n", err)
				// 抓取失败也记录到0条数据中
				recordZeroDataCode(code.Code)
			}

			// 控制请求间隔
			time.Sleep(500 * time.Millisecond)
		}(indexCode)
	}

	// 等待所有任务完成
	wg.Wait()
	fmt.Println("所有指数数据抓取完成，开始统计0条数据的股票代码")

	// 将0条数据的股票代码写入文件
	writeZeroDataCodesToFile()
	fmt.Printf("已将%d个0条数据的股票代码写入zero_data_codes.txt\n", len(zeroDataCodes))
}

// FetchData 抓取指数数据
func FetchData(indexCode IndexCode, client *resty.Client) error {
	fmt.Printf("fetch %s date -> begin\n", indexCode.Code)

	// 构建请求URL
	url := fmt.Sprintf(
		"https://push2his.eastmoney.com/api/qt/stock/kline/get?secid=%s&fields1=f1%%2Cf2%%2Cf3%%2Cf4%%2Cf5%%2Cf6&fields2=f51%%2Cf52%%2Cf53%%2Cf54%%2Cf55%%2Cf56%%2Cf57%%2Cf58%%2Cf59%%2Cf60%%2Cf61&klt=101&fqt=1&beg=0&end=20500101&lmt=120",
		indexCode.Secid,
	)

	// 发送请求
	response, err := client.R().
		Get(url)
	if err != nil {
		return fmt.Errorf("发送请求失败: %v", err)
	}

	// 检查响应状态码
	if response.StatusCode() != http.StatusOK {
		return fmt.Errorf("请求返回错误状态码: %d, 响应内容: %s",
			response.StatusCode(), string(response.Body()))
	}

	// 处理响应编码
	body, err := detectAndConvertEncoding(response.Body())
	if err != nil {
		return fmt.Errorf("处理响应编码失败: %v", err)
	}

	// 解析响应
	var eastmoneyResponse EastmoneyResponse
	if err := json.Unmarshal(body, &eastmoneyResponse); err != nil {
		// 输出部分响应内容用于调试
		preview := string(body)
		if len(preview) > 200 {
			preview = preview[:200] + "..."
		}
		return fmt.Errorf("解析JSON失败: %v, 响应预览: %s", err, preview)
	}

	// 处理K线数据
	var indexDataList []IndexData
	var wg sync.WaitGroup
	errorCount := 0
	errorMutex := &sync.Mutex{}

	// 创建数据通道
	dataChan := make(chan IndexData, len(eastmoneyResponse.Data.Klines))

	// 并行处理K线数据
	for _, item := range eastmoneyResponse.Data.Klines {
		wg.Add(1)
		go func(kline string) {
			defer wg.Done()

			itemSplitVec := strings.Split(kline, ",")
			if len(itemSplitVec) < 11 {
				errorMutex.Lock()
				errorCount++
				errorMutex.Unlock()
				return
			}

			date := itemSplitVec[0]
			closePoint, err := strconv.ParseFloat(itemSplitVec[2], 64)
			if err != nil {
				errorMutex.Lock()
				errorCount++
				errorMutex.Unlock()
				fmt.Printf("解析收盘价失败: %v, 数据: %s\n", err, kline)
				return
			}

			indexData := IndexData{
				Date:       date,
				ClosePoint: closePoint,
			}
			dataChan <- indexData
		}(item)
	}

	// 等待所有 goroutine 完成并关闭通道
	go func() {
		wg.Wait()
		close(dataChan)
	}()

	// 从通道收集数据
	for data := range dataChan {
		indexDataList = append(indexDataList, data)
	}

	if errorCount > 0 {
		fmt.Printf("处理 %s 数据时共有 %d 条记录解析失败\n", indexCode.Code, errorCount)
	}

	// 创建目录（如果不存在）
	if err := os.MkdirAll("index-data", 0755); err != nil {
		return fmt.Errorf("创建目录失败: %v", err)
	}

	// 写入JSON文件
	filePath := filepath.Join("index-data", fmt.Sprintf("%s.json", eastmoneyResponse.Data.Code))
	jsonData, err := json.MarshalIndent(indexDataList, "", "  ")
	if err != nil {
		return fmt.Errorf("JSON序列化失败: %v", err)
	}

	if err := os.WriteFile(filePath, jsonData, 0644); err != nil {
		return fmt.Errorf("写入文件失败: %v", err)
	}

	fmt.Printf("fetch %s date <- end, 共 %d 条数据\n", indexCode.Code, len(indexDataList))

	// 记录0条数据的股票代码
	if len(indexDataList) == 0 {
		recordZeroDataCode(indexCode.Code)
	}

	return nil
}

// detectAndConvertEncoding 检测并转换响应编码
func detectAndConvertEncoding(body []byte) ([]byte, error) {
	// 尝试检测编码
	reader, err := charset.NewReader(bytes.NewReader(body), "")
	if err != nil {
		return body, nil // 无法检测，返回原始内容
	}

	// 读取转换后的内容
	var buf bytes.Buffer
	_, err = io.Copy(&buf, reader)
	if err != nil {
		return body, err
	}

	return buf.Bytes(), nil
}

// recordZeroDataCode 记录0条数据的股票代码
func recordZeroDataCode(code string) {
	zeroDataMutex.Lock()
	defer zeroDataMutex.Unlock()
	zeroDataCodes = append(zeroDataCodes, code)
}

// writeZeroDataCodesToFile 将0条数据的股票代码写入文件
func writeZeroDataCodesToFile() {
	file, err := os.Create("zero_data_codes.txt")
	if err != nil {
		log.Fatalf("创建文件失败: %v", err)
	}
	defer file.Close()

	// 去重
	codeMap := make(map[string]bool)
	var uniqueCodes []string
	for _, code := range zeroDataCodes {
		if !codeMap[code] {
			codeMap[code] = true
			uniqueCodes = append(uniqueCodes, code)
		}
	}

	// 按行写入股票代码
	for _, code := range uniqueCodes {
		_, err := file.WriteString(code + "\n")
		if err != nil {
			log.Printf("写入代码 %s 失败: %v", code, err)
		}
	}
}
